# coding=utf8

from __future__ import print_function
import tensorflow as tf
import tflearn
import tensorflow.contrib.rnn as rnn
mnist = tflearn.datasets.mnist.read_data_sets(train_dir='/Users/vista/PycharmProjects/data/MNIST_data',one_hot=True)

# 一张图片是28*28,FNN是一次性把数据输入到网络，RNN把它分成块
chunk_size = 28
chunk_n = 28

rnn_size = 256

n_output_layer = 10  # 输出层

X = tf.placeholder('float', [None, chunk_n, chunk_size])
Y = tf.placeholder("float", shape=[None, n_output_layer])

batch_size = 50


# 定义待训练的神经网络
def recurrent_neural_network(data):
    layer = {'w_': tf.Variable(tf.random_normal([rnn_size, n_output_layer])),
             'b_': tf.Variable(tf.random_normal([n_output_layer]))}

    lstm_cell = rnn.BasicLSTMCell(rnn_size)
    #lstm_cell = tf_vista.nn.rnn_cell.BasicLSTMCell(rnn_size)

    # batch * n * size ==> n * batch * size
    # [1,0,2] means 0 <=> 1 维互换
    data = tf.transpose(data, [1, 0, 2])

    data = tf.reshape(data, [-1, chunk_size])
    data = tf.split(axis=0,num_or_size_splits=chunk_n, value=data)
    outputs, status = rnn.static_rnn(cell=lstm_cell, inputs = data, dtype=tf.float32)

    ouput = tf.add(tf.matmul(outputs[-1], layer['w_']), layer['b_'])

    return ouput


# 使用数据训练神经网络
def train_neural_network(X, Y):
    py_x = recurrent_neural_network(X)
    cost_func = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits( logits= py_x, labels=Y))
    optimizer = tf.train.AdamOptimizer(0.01).minimize(cost_func)

    epochs = 13
    with tf.Session() as session:
        session.run(tf.global_variables_initializer())

        for epoch in range(epochs):

            epoch_loss = 0

            for i in range(int(mnist.train.num_examples / batch_size)):
                x, y = mnist.train.next_batch(batch_size)
                x = x.reshape([batch_size, chunk_n, chunk_size])
                _, c = session.run([optimizer, cost_func], feed_dict={X: x, Y: y})
                epoch_loss += c
                if i % 100 == 0:
                    print (c)

            print(epoch, ' : ', epoch_loss)

        correct = tf.equal(tf.argmax(py_x, 1), tf.argmax(Y, 1))
        accuracy = tf.reduce_mean(tf.cast(correct, 'float'))
        print('accuracy: ', accuracy.eval({X: mnist.test.images.reshape(-1, chunk_n, chunk_size), Y: mnist.test.labels}))


train_neural_network(X, Y)
